from torch.utils.data import Dataset
import torch
import numpy as np


class Net_Dataset(Dataset):
    def __init__(self, x_data, y_data):
        self.x_data = x_data
        self.y_data = np.array(y_data, dtype=np.longlong)
        print(self.x_data.shape)

    #根据索引index返回dataset[index]
    def __getitem__(self, index):
        # x = torch.from_numpy(self.x_data[index]) # (samples, 1, 256, 256)  # samples表示batch_size
        x = self.x_data[index]
        y = self.y_data[index]
        return x, y

# 返回整个数据集大小
    def __len__(self):
        return self.y_data.shape[0]
